package _01_第一季.树04.Tree01;

import java.util.Comparator;

public class RBTree<E> extends BBST<E> {

    private static final boolean RED = false;
    private static final boolean BLACK = true;

    public RBTree(){
        this(null);
    }
    public RBTree(Comparator<E> comparator){
        super(comparator);
    }

    private static class RBNode<E> extends Node<E>{
        boolean color = RED;
        public RBNode(){
            this(null,null);
        }
        public RBNode(E element,Node<E> parent){
            super(element,parent);
        }
        @Override
        public String toString() {
            String str = "";
            if(color == RED){
                str = "R_";
            }
            return str+element.toString();
        }
    }

    @Override
    protected void afterAdd(Node<E> node) {
        Node<E> parent = node.parent;
        if(parent == null){
            black(node);
            return;
        }

        if(isBlack(parent)) return;

        Node<E> uncle = parent.sibling();
        Node<E> grand = parent.parent;

        if(isRed(uncle)){
            black(parent);
            black(uncle);
            afterAdd(red(grand));
            return;
        }

        if(parent.isLeftChild()){
            red(grand);
            if(node.isLeftChild()){ //LL
                black(parent);
            }else{ //LR
                black(node);
                rotateLeft(parent);
            }
            rotateRight(grand);
        }else{
            red(grand);
            if(node.isLeftChild()){ //RL
                black(node);
                rotateRight(parent);
            }else{ //RR
                black(parent);
            }
            rotateLeft(grand);
        }

    }

    @Override
    protected void afterRemove(Node<E> node) {
       if(isRed(node)){
           black(node);
           return;
       }

       Node<E> parent = node.parent;
       if(parent == null) return;

       boolean left = parent.left == null || node.isLeftChild();
       Node<E> sibling = left ? parent.right : parent.left;

       if(left){
           if(isRed(sibling)){
               black(sibling);
               red(parent);
               rotateLeft(parent);
               sibling = parent.right;
           }

           if(isBlack(sibling.left) && isBlack(sibling.right)){
               boolean parentBlack = isBlack(parent);
               black(parent);
               red(sibling);
               if(parentBlack){
                   afterRemove(parent);
               }
           }else{
               if(isBlack(sibling.right)){
                   rotateRight(sibling);
                   sibling = parent.right;
               }
               color(sibling,colorOf(parent));
               black(parent);
               black(sibling.left);
               rotateLeft(parent);
           }
       } else{
           if(isRed(sibling)){
               black(sibling);
               red(parent);
               rotateRight(parent);
               sibling = parent.left;
           }

           if(isBlack(sibling.left) && isBlack(sibling.right)){
               boolean parentBlack = isBlack(parent);
               black(parent);
               red(sibling);
               if(parentBlack){
                   afterRemove(parent);
               }
           }else{
                if(isBlack(sibling.left)){
                    rotateLeft(sibling);
                    sibling = parent.left;
                }
                color(sibling,colorOf(parent));
                black(parent);
                black(sibling.left);
                rotateRight(parent);
           }
       }
    }

    private Node<E> color(Node<E> node, boolean color){
        if(node == null) return null;
        ((RBNode<E>)node).color = color;
        return node;
    }

    private Node<E> red(Node<E> node){
        return color(node,RED);
    }

    private Node<E> black(Node<E> node){
        return color(node,BLACK);
    }

    private boolean colorOf(Node<E> node){
        return node == null ? BLACK : ((RBNode<E>)node).color;
    }

    private boolean isBlack(Node<E> node){
        return colorOf(node) == BLACK;
    }

    private boolean isRed(Node<E> node){
        return colorOf(node) == RED;
    }

    @Override
    protected Node<E> createNode(E element, Node<E> parent) {
        return new RBNode<>(element,parent);
    }
}
